package main

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
	"github.com/volcengine/volcengine-go-sdk/volcengine"
	"log"
	"net/http"
	"os"
)

/*func main() {
	client := arkruntime.NewClientWithApiKey(
		os.Getenv("ARK_API_KEY"),
	)
	ctx := context.Background()

	fmt.Println("----- multiple rounds request -----")
	req := model.BotChatCompletionRequest{
		// 将<YOUR_BOT_ID>替换为您的应用ID
		Model: "bot-20250425164816-2p259",
		Messages: []*model.ChatCompletionMessage{
			{
				Role: model.ChatMessageRoleSystem,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String("你是豆包，字节跳动开发的 AI 人工智能助手"),
				},
			},
			{
				Role: model.ChatMessageRoleUser,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String("花椰菜是什么？"),
				},
			},
			{
				Role: model.ChatMessageRoleAssistant,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String("花椰菜又称菜花、花菜，是一种常见的蔬菜。"),
				},
			},
			{
				Role: model.ChatMessageRoleUser,
				Content: &model.ChatCompletionMessageContent{
					StringValue: volcengine.String("再详细点"),
				},
			},
		},
	}

	resp, err := client.CreateBotChatCompletion(ctx, req)
	if err != nil {
		fmt.Printf("standard chat error: %v\n", err)
		return
	}
	fmt.Println(*resp.Choices[0].Message.Content.StringValue)
	if resp.References != nil {
		for _, ref := range resp.References {
			fmt.Printf("reference url: %s\n", ref.Url)
		}
	}
}*/

// ChatClient 封装对话客户端和相关配置
type ChatClient struct {
	client       *arkruntime.Client
	botID        string
	systemPrompt string
}

// NewChatClient 创建新的对话客户端
func NewChatClient() *ChatClient {
	return &ChatClient{
		client:       arkruntime.NewClientWithApiKey(os.Getenv("ARK_API_KEY")),
		botID:        "bot-20250425164816-2p259",
		systemPrompt: "你是豆包，字节跳动开发的 AI 人工智能助手",
	}
}

// ChatRequest 对话请求结构
type ChatRequest struct {
	Messages []*model.ChatCompletionMessage `json:"messages"`
}

// ChatResponse 对话响应结构
type ChatResponse struct {
	Content    string   `json:"content"`
	References []string `json:"references,omitempty"`
}

// ChatCompletion 核心对话接口函数
func (c *ChatClient) ChatCompletion(ctx context.Context, req ChatRequest) (*ChatResponse, error) {
	// 自动添加系统提示
	messages := append([]*model.ChatCompletionMessage{
		{
			Role: model.ChatMessageRoleSystem,
			Content: &model.ChatCompletionMessageContent{
				StringValue: volcengine.String(c.systemPrompt),
			},
		},
	}, req.Messages...)

	apiReq := model.BotChatCompletionRequest{
		Model:    c.botID,
		Messages: messages,
	}

	resp, err := c.client.CreateBotChatCompletion(ctx, apiReq)
	if err != nil {
		return nil, err
	}

	response := &ChatResponse{
		Content: *resp.Choices[0].Message.Content.StringValue,
	}

	if resp.References != nil {
		response.References = make([]string, 0, len(resp.References))
		for _, ref := range resp.References {
			response.References = append(response.References, ref.Url)
		}
	}

	return response, nil
}

// 在main函数前添加CORS中间件
func enableCORS(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 设置允许的源（生产环境应替换为具体域名）
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
		w.Header().Set("Access-Control-Allow-Headers", "Accept, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization")

		// 处理预检请求
		if r.Method == "OPTIONS" {
			w.WriteHeader(http.StatusNoContent)
			return
		}

		next.ServeHTTP(w, r)
	})
}

// HTTP 处理函数示例
func main() {
	chatClient := NewChatClient()
	mux := http.NewServeMux()
	mux.HandleFunc("/chat", func(w http.ResponseWriter, r *http.Request) {
		var req ChatRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		resp, err := chatClient.ChatCompletion(r.Context(), req)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(resp)
	})

	// 应用CORS中间件
	handler := enableCORS(mux)
	// 启动服务
	fmt.Println("Server running on :8080")
	log.Fatal(http.ListenAndServe(":8080", handler))
}
